// Project partly adapted from https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf

#include <iostream>
#include <math.h>

// This will output the proper CUDA error strings in the event
// that a CUDA host call returns an error
#define checkGPUErrors(val) __check((val), #val, __FILE__, __LINE__)

void __check(cudaError_t result, char const *const func, const char *const file,
           int const line) {
  if (result) {
    fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\" \n", file, line,
            static_cast<unsigned int>(result), cudaGetErrorName(result), func);
    exit(EXIT_FAILURE);
  }
}

const int NUM_ELEMENT = 32*1024*1024; // 32 M elements
const int THREAD_PER_BLOCK = 256;

// 每个block负责将THREAD_PER_BLOCK个数据累加，结果写到z[blockIdx.x]
__global__ void reduce(float *x, float *z)
{
    // threadIdx.x是线程在block内的索引(0 到 THREAD_PER_BLOCK-1)
    int tid = threadIdx.x;
    
    // blockIdx.x是当前block的索引(0 到 grid_dim-1)
    int bid = blockIdx.x;
    
    // idx计算线程对应的全局内存位置:
    // - bid * blockDim.x 计算当前block在全局内存中的起始位置
    // - 加上tid得到当前线程对应的具体位置
    // 例如：假设THREAD_PER_BLOCK=256
    // - block 0的线程访问 0-255
    // - block 1的线程访问 256-511
    // - block 2的线程访问 512-767
    int idx = bid * blockDim.x + tid;
    
    // 在每个block内进行归约
    for (int s = 1; s < THREAD_PER_BLOCK; s *= 2) {
        // tid % (2*s) == 0 选择执行归约的线程
        // 例如：第一轮s=1时，线程0,2,4,6...执行
        //      第二轮s=2时，线程0,4,8,12...执行
        //      第三轮s=4时，线程0,8,16,24...执行
        // (tid + s) < THREAD_PER_BLOCK 确保不会访问超出block范围的数据
        if (tid % (2*s) == 0 && (tid + s) < THREAD_PER_BLOCK) {
            // 在全局内存中:
            // x[idx] 是当前线程的数据位置
            // x[idx+s] 是需要归约的相邻数据位置
            x[idx] = x[idx] + x[idx+s];
        }
        
        // 确保block内所有线程完成当前步骤
        __syncthreads();
    }
    
    // 每个block只需要一个线程写回结果
    // 选择block内的第一个线程(tid=0)
    // 将最终结果写入z数组对应位置
    if(tid == 0) {
        z[bid] = x[idx];  // idx此时是block起始位置
    }
}


int main(void)
{
    int array_size = NUM_ELEMENT; 
    int z_size = array_size/THREAD_PER_BLOCK;

    float *x, *z;
    x = (float*)malloc(array_size*sizeof(float));
    z = (float*)malloc(z_size*sizeof(float));

    float *x_gpu, *z_gpu; 
    checkGPUErrors(cudaMalloc(&x_gpu, array_size*sizeof(float)));
    checkGPUErrors(cudaMalloc(&z_gpu, z_size*sizeof(float)));

    // initialize x arrays on the host
    for (int i = 0; i < array_size; i++) {
        x[i] = 1.0f;
    }
    checkGPUErrors(cudaMemcpy(x_gpu, x, array_size*sizeof(float), cudaMemcpyHostToDevice));

    cudaEvent_t gpu_start, gpu_stop;
    cudaEventCreate(&gpu_start);
    cudaEventCreate(&gpu_stop);
    cudaEventRecord(gpu_start);
    int block_dim = THREAD_PER_BLOCK;
    int grid_dim = array_size / THREAD_PER_BLOCK;
    reduce<<<grid_dim,block_dim>>>(x_gpu, z_gpu);
    // 等待GPU上的kernel执行完毕
    cudaEventRecord(gpu_stop);
    cudaEventSynchronize(gpu_stop);  
    float milliseconds = 0;
    cudaEventElapsedTime(&milliseconds, gpu_start, gpu_stop);

    printf("Execution time measured on GPU: %f us.\n", milliseconds*1000);

    //检验是否所有的结果都是THREAD_PER_BLOCK
    checkGPUErrors(cudaMemcpy(z, z_gpu, z_size*sizeof(float), cudaMemcpyDeviceToHost));
    float maxError = 0.0f;
    for(int i = 0; i<z_size; i++) {
      maxError = fmax(maxError, fabs(z[i] - THREAD_PER_BLOCK));
    }
    if(maxError < 1e-6)
        std::cout << "PASS. MaxError: " << maxError << std::endl;
    else
        std::cout << "FAIL. MaxError: " << maxError << std::endl;

    // 计算访存比
    // 对于1024个数字的归约求和，需要1023次乘加
    // 对于THREAD_PER_BLOCK个数字的归约求和，需要THREAD_PER_BLOCK-1次加法
    float gflops = (float)(THREAD_PER_BLOCK - 1) * grid_dim / 1024 / 1024 / milliseconds;
    printf("Throughput: %f GFLOP/s.\n", gflops);

    // Free memory
    checkGPUErrors(cudaFree(x_gpu));
    checkGPUErrors(cudaFree(z_gpu));
    free(x);
    free(z);

    return 0;
}
